#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0] == 'extract_edgels-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io
from scipy.optimize import fmin, fmin_powell

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from quaternion import Quat

import simplejson

import Image

set_printoptions(precision=3)


def create_pic(job_params, framenum):
    filename = job_params['root_directory']+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(job_params["gaussian_smoothing_factor"]))


    ## Creates picture object
    if not 'projection_model' in job_params:
        model = 'perspective'
    else:
        model = job_params['projection_model']

    if model == 'perspective':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance_guess']
        p_point = array(job_params['principal_point_guess'])

        pic = PicturePinhole(imr, focal_distance, p_point[0], p_point[1])

    elif model == 'equirectangular':
        pic = PictureEqrec(imr)

    return pic

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0] == 'extract_edgels-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load images and initialize pic objects

    ## Sets filename from input argument
    if len(sys.argv) < 2:
        print sys.argv[0], '<job_file.json>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    picL = []

    ## picz
    for framenum in range(24):
        #framenum = int(sys.argv[2])
        picL.append(create_pic(job_params, framenum))

    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold']

    for pic in picL:
        ## Perform initial orientation estimation
        pic.extract_edgels(gspc, glim)
        pic.generate_random_edgels()


    for pic in picL:
        pic.middlex,pic.middley = job_params['principal_point_guess']
        pic.fd = job_params['focal_distance_guess']

        ## Re-calculate all normals with the current intrinsic parameters
        pic.calculate_edgel_normals()
        pic.calculate_random_normals()

        ## First orientation estimation, using L1 + tukey.
        #xini=[]
        #xini = [array([0.949244, -0.184524, 0.254289, 0.014982])]
        xini = array(Quat(x).q for x in job_params['initial_guesses'])
        xopt, sig_hat = pic.find_orientation(xini, 0.0)
        pic.orientation = xopt.canonical()
        # pic.orientation, sig_hat = Quat(0.934345, -0.031462, 0.354394, -0.020342), .15
        print 'Ori est', pic.orientation

        pic.sigma = sig_hat
        pic.set_labels(pic.sigma)

    plot_aux.plot_problem(picL[0])


    def target(x, pic,fp):
        q = Quat(x).normalize()
        return pic.angle_error(pic.middlex, pic.middley, pic.fd,
                        (pic.orientation*q).q, fp)
    def target2(x, pic,fp):
        orL = [Quat(array([x[k],x[k+1],x[k+2]])) for k in range(3, len(x), 3)]
        tot = 0
        for (ori,pic) in zip(orL,picL):
            tot += pic.angle_error(x[0],x[1],x[2],
                                   (pic.orientation*ori).q, fp)
        return tot

    pic = picL[0]
    x0 = [0.,0.,0.]
    fp = (1.,1,0)
    for pic in picL:
        sol = fmin(target, x0, (pic,fp))
        pic.orientation = pic.orientation*Quat(sol)
        print sol

    ##################################################################
    ## Optimization on the intrinsic parameters, and on all camera
    ## orientations at the same time.
    fp = (1/.15, 3, 1.5)
    x0 = (pic.middlex, pic.middley, pic.fd)+ len(picL)*(0,0,0)
    x = fmin(target2, x0, (pic,fp))
    ##
    ##################################################################

    ## Put results into the picture objects
    ip_opt = x[:3]
    orL = [Quat(array([x[k],x[k+1],x[k+2]])) for k in range(3, len(x), 3)]

    for (ori,pic) in zip(orL,picL):
        pic.old_orientation = pic.orientation
        pic.orientation = pic.orientation*ori
        pic.middlex = ip_opt[0]
        pic.middley = ip_opt[1]
        pic.fd = ip_opt[2]
        ## Get new normals to plot the sphere...
        pic.calculate_edgel_normals()
        pic.set_labels(pic.sigma)

    pic=picL[0]
    rez=pic.angle_error_res(pic.middlex, pic.middley, pic.fd,
                        pic.orientation.q, fp)

    

    plot_aux.plot_problem(picL[0],10)

